import timm
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import numpy as np
class ResNet(nn.Module):
    def __init__(self):
        super(ResNet, self).__init__()
        self.m = timm.create_model('resnet50', pretrained=True, features_only=True,out_indices=(2,3,4))

    def forward(self,x):
        #新增y变量,用于可视化倒残差模块的处理结果
        #y = x

        x = self.m(x)
        output = x
        return tuple(output)

if __name__ == '__main__':
    net = ResNet()
    x = torch.randn(2, 3, 320, 320)
    o = net(x)
    for i in o:
        print(i.shape)